"""
向量数据库管理
基于ChromaDB的向量存储和检索
"""

import os
import logging
import json
from typing import List, Dict, Any, Optional
import chromadb
from chromadb.config import Settings
from .embeddings import Embeddings

logger = logging.getLogger(__name__)

class VectorStore:
    """向量数据库管理类"""
    
    def __init__(self, persist_directory: str = "./chroma_db"):
        """
        初始化向量数据库
        
        Args:
            persist_directory: 向量数据库持久化目录
        """
        self.persist_directory = persist_directory
        self.embeddings = Embeddings()
        
        # 确保目录存在
        os.makedirs(persist_directory, exist_ok=True)
        
        # 初始化ChromaDB客户端
        self.client = chromadb.PersistentClient(
            path=persist_directory,
            settings=Settings(
                anonymized_telemetry=False,
                allow_reset=True
            )
        )
        
        # 初始化集合
        self._init_collections()
        
        logger.info(f"向量数据库初始化成功，存储路径: {persist_directory}")
    
    def _init_collections(self):
        """初始化数据库集合"""
        # 简历集合
        self.resume_collection = self.client.get_or_create_collection(
            name="resumes",
            metadata={"description": "简历向量存储"}
        )
        
        # 职位描述集合
        self.job_collection = self.client.get_or_create_collection(
            name="jobs",
            metadata={"description": "职位描述向量存储"}
        )
        
        # 技能库集合
        self.skill_collection = self.client.get_or_create_collection(
            name="skills",
            metadata={"description": "技能库向量存储"}
        )
        
        # 面试题库集合
        self.interview_collection = self.client.get_or_create_collection(
            name="interviews",
            metadata={"description": "面试题库向量存储"}
        )
        
        logger.info("数据库集合初始化完成")
    
    def add_resume(self, resume_id: str, content: str, metadata: Dict[str, Any] = None):
        """
        添加简历到向量数据库
        
        Args:
            resume_id: 简历ID
            content: 简历内容
            metadata: 元数据
        """
        try:
            # 生成向量
            embedding = self.embeddings.embed_text(content)
            
            # 存储到数据库
            self.resume_collection.add(
                embeddings=[embedding],
                documents=[content],
                metadatas=[metadata or {}],
                ids=[resume_id]
            )
            
            logger.info(f"简历 {resume_id} 已添加到向量数据库")
            
        except Exception as e:
            logger.error(f"添加简历失败: {str(e)}")
            raise
    
    def add_job(self, job_id: str, title: str, description: str, metadata: Dict[str, Any] = None):
        """
        添加职位描述到向量数据库
        
        Args:
            job_id: 职位ID
            title: 职位标题
            description: 职位描述
            metadata: 元数据
        """
        try:
            # 组合标题和描述
            full_content = f"职位: {title}\n描述: {description}"
            
            # 生成向量
            embedding = self.embeddings.embed_text(full_content)
            
            # 存储到数据库
            self.job_collection.add(
                embeddings=[embedding],
                documents=[full_content],
                metadatas=[metadata or {}],
                ids=[job_id]
            )
            
            logger.info(f"职位 {job_id} 已添加到向量数据库")
            
        except Exception as e:
            logger.error(f"添加职位失败: {str(e)}")
            raise
    
    def add_skill(self, skill_id: str, skill_name: str, description: str, metadata: Dict[str, Any] = None):
        """
        添加技能到向量数据库
        
        Args:
            skill_id: 技能ID
            skill_name: 技能名称
            description: 技能描述
            metadata: 元数据
        """
        try:
            # 组合技能信息
            full_content = f"技能: {skill_name}\n描述: {description}"
            
            # 生成向量
            embedding = self.embeddings.embed_text(full_content)
            
            # 存储到数据库
            self.skill_collection.add(
                embeddings=[embedding],
                documents=[full_content],
                metadatas=[metadata or {}],
                ids=[skill_id]
            )
            
            logger.info(f"技能 {skill_id} 已添加到向量数据库")
            
        except Exception as e:
            logger.error(f"添加技能失败: {str(e)}")
            raise
    
    def search_resumes(self, query: str, n_results: int = 5) -> List[Dict[str, Any]]:
        """
        搜索简历
        
        Args:
            query: 搜索查询
            n_results: 返回结果数量
            
        Returns:
            搜索结果列表
        """
        try:
            # 生成查询向量
            query_embedding = self.embeddings.embed_text(query)
            
            # 搜索
            results = self.resume_collection.query(
                query_embeddings=[query_embedding],
                n_results=n_results
            )
            
            # 格式化结果
            formatted_results = []
            for i in range(len(results['ids'][0])):
                formatted_results.append({
                    'id': results['ids'][0][i],
                    'content': results['documents'][0][i],
                    'metadata': results['metadatas'][0][i],
                    'distance': results['distances'][0][i]
                })
            
            return formatted_results
            
        except Exception as e:
            logger.error(f"搜索简历失败: {str(e)}")
            raise
    
    def search_jobs(self, query: str, n_results: int = 5) -> List[Dict[str, Any]]:
        """
        搜索职位
        
        Args:
            query: 搜索查询
            n_results: 返回结果数量
            
        Returns:
            搜索结果列表
        """
        try:
            # 生成查询向量
            query_embedding = self.embeddings.embed_text(query)
            
            # 搜索
            results = self.job_collection.query(
                query_embeddings=[query_embedding],
                n_results=n_results
            )
            
            # 格式化结果
            formatted_results = []
            for i in range(len(results['ids'][0])):
                formatted_results.append({
                    'id': results['ids'][0][i],
                    'content': results['documents'][0][i],
                    'metadata': results['metadatas'][0][i],
                    'distance': results['distances'][0][i]
                })
            
            return formatted_results
            
        except Exception as e:
            logger.error(f"搜索职位失败: {str(e)}")
            raise
    
    def search_skills(self, query: str, n_results: int = 5) -> List[Dict[str, Any]]:
        """
        搜索技能
        
        Args:
            query: 搜索查询
            n_results: 返回结果数量
            
        Returns:
            搜索结果列表
        """
        try:
            # 生成查询向量
            query_embedding = self.embeddings.embed_text(query)
            
            # 搜索
            results = self.skill_collection.query(
                query_embeddings=[query_embedding],
                n_results=n_results
            )
            
            # 格式化结果
            formatted_results = []
            for i in range(len(results['ids'][0])):
                formatted_results.append({
                    'id': results['ids'][0][i],
                    'content': results['documents'][0][i],
                    'metadata': results['metadatas'][0][i],
                    'distance': results['distances'][0][i]
                })
            
            return formatted_results
            
        except Exception as e:
            logger.error(f"搜索技能失败: {str(e)}")
            raise
    
    def get_collection_stats(self) -> Dict[str, Any]:
        """获取数据库统计信息"""
        try:
            stats = {}
            
            # 获取各集合的文档数量
            stats['resumes'] = self.resume_collection.count()
            stats['jobs'] = self.job_collection.count()
            stats['skills'] = self.skill_collection.count()
            stats['interviews'] = self.interview_collection.count()
            
            return stats
            
        except Exception as e:
            logger.error(f"获取统计信息失败: {str(e)}")
            raise
    
    def delete_resume(self, resume_id: str):
        """删除简历"""
        try:
            self.resume_collection.delete(ids=[resume_id])
            logger.info(f"简历 {resume_id} 已删除")
        except Exception as e:
            logger.error(f"删除简历失败: {str(e)}")
            raise
    
    def delete_job(self, job_id: str):
        """删除职位"""
        try:
            self.job_collection.delete(ids=[job_id])
            logger.info(f"职位 {job_id} 已删除")
        except Exception as e:
            logger.error(f"删除职位失败: {str(e)}")
            raise
